#
# Authors: Simon Vandenhende
# Licensed under the CC BY-NC 4.0 license (https://creativecommons.org/licenses/by-nc/4.0/)

import os
import cv2
import imageio
import numpy as np
import json
import torch
import scipy.io as sio
from utils.utils import get_output, mkdir_if_missing
from evaluation.eval_meters import HeightMeter, SemsegMeter
from termcolor import colored


class PerformanceMeter(object):
    """ A general performance meter which shows performance across one or more tasks """
    def __init__(self, p):
        self.database = p['train_db_name']
        self.tasks = p.TASKS.NAMES
        self.meters = {t: get_single_task_meter(p, self.database, t) for t in self.tasks}

    def reset(self):
        for t in self.tasks:
            self.meters[t].reset()

    def update(self, pred, gt):
        for t in self.tasks:
            self.meters[t].update(pred[t].detach().cpu().numpy(), gt[t].detach().cpu().numpy())

    def get_score(self, verbose=True):
        eval_dict = {}
        for t in self.tasks:
            eval_dict[t] = self.meters[t].get_score(verbose)

        return eval_dict


def get_single_task_meter(p, database, task):
    """ Retrieve a meter to measure the single-task performance """
    if task == 'semseg':
        return SemsegMeter(database)

    elif task == 'height':
        return HeightMeter()

    else:
        raise NotImplementedError


@torch.no_grad()
def eval_result(p, val_loader, model):
    model.eval()

    results = {}

    height_dic = HeightMeter()
    seg_dic = SemsegMeter(p['train_db_name'])

    for ii, sample in enumerate(val_loader):
        inputs, meta = sample['image'].cuda(non_blocking=True), sample['meta']
        output = model(inputs)

        for task in p.TASKS.NAMES:
            pred = get_output(output[task], task).detach().cpu().data.numpy()
            label = sample[task].detach().cpu().data.numpy()
            if task == 'height':
                if p['heightloss'] == 'ssimse':
                    mask = label > 0
                    pred = align_height(pred, label, mask)
                height_dic.update(pred, label)

            elif task == 'semseg':
                seg_dic.update(pred, label)

    if 'height' in p.TASKS.NAMES:
        results['height'] = height_dic.get_score(is_val=True)

    if 'semseg' in p.TASKS.NAMES:
        results['semseg'] = seg_dic.get_score(is_val=True)

    return results


def compute_scale_and_shift(prediction, target, mask):
    # system matrix: A = [[a_00, a_01], [a_10, a_11]]
    a_00 = np.sum(mask * prediction * prediction, (1, 2))
    a_01 = np.sum(mask * prediction, (1, 2))
    a_11 = np.sum(mask, (1, 2))

    # right hand side: b = [b_0, b_1]
    b_0 = np.sum(mask * prediction * target, (1, 2))
    b_1 = np.sum(mask * target, (1, 2))

    # solution: x = A^-1 . b = [[a_11, -a_01], [-a_10, a_00]] / (a_00 * a_11 - a_01 * a_10) . b
    x_0 = np.zeros_like(b_0)
    x_1 = np.zeros_like(b_1)

    det = a_00 * a_11 - a_01 * a_01
    # A needs to be a positive definite matrix.
    valid = det > 0

    x_0[valid] = (a_11[valid] * b_0[valid] - a_01[valid] * b_1[valid]) / det[valid]
    x_1[valid] = (-a_01[valid] * b_0[valid] + a_00[valid] * b_1[valid]) / det[valid]

    return x_0, x_1


def align_height(prediction, target, mask):
    # transform predicted disparity to aligned height
    prediction = prediction.squeeze(1)
    target = target.squeeze(1)
    mask = mask.squeeze(1)

    target_norm = np.zeros_like(target)
    target_norm[mask == 1] = (target[mask == 1] - target[mask == 1].min()) / (target[mask == 1].max() - target[mask == 1].min())

    scale, shift = compute_scale_and_shift(prediction, target_norm, mask)
    prediction_aligned = scale.reshape((-1, 1, 1)) * prediction + shift.reshape((-1, 1, 1))

    prediciton_height = prediction_aligned*(target[mask == 1].max() - target[mask == 1].min()) + target[mask == 1].min()

    return np.expand_dims(prediciton_height, axis=1)